#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <iostream>
#include <sstream>

// MACA runtime
#include <mcblas/mcblas.h>
#include <mc_runtime.h>
#include "helper.h"

const int kNRepeats = 5; // Number of repeats for timing. Can be overwrite with -t option
const int kStart = 1024;
const int kStop = 4096;
const int kStep = 256;

void print_help(const char* program_name) {
    std::stringstream help;
    help << "Usage: " << program_name << " [options] [M N K]"
         << "\nOptions:"
         << "\n\t-h\tShow this help message"
         << "\n\t-t\tTrace mode: Execute the tested function once"
         << "\n"
         << "\nUse Cases:"
         << "\n\tNo arguments\t\tPerform a sweep of square matrix multiplications over the range (1024, 4096, 256)"
         << "\n\t-t\t\t\tPerform a trace and sweep of square matrix multiplications over the range (1024, 4096, 256)"
         << "\n\tM N K\t\t\tTest a specific GEMM problem size with matrices of dimensions M x N and N x K"
         << "\n\t-t M N K\t\tPerform a trace of the specific GEMM problem size with matrices of dimensions M x N and N x K"
         << "\n";
    std::cout << help.str();
}


void SGEMM_GPU_REF(mcblasHandle_t, int, int, int, float *, float *, float *);
void SGEMM_GPU(mcblasHandle_t, int, int, int, float *, float *, float *);
void random_matrix(int, int, float *);
float compare_matrices(int, int, float *, float *);

// This function is used to call the kernel, compare with mcBLAS, and measure the performance
void callKernel(int m, int n, int k, bool trace = false) {
  double diff = 0;
  int num_rep = trace ? 1 : kNRepeats;
  float *a, *b, *c, *cref, *cold;

  mcblasHandle_t handle;
  checkGPUErrors(mcblasCreate(&handle));
  // checkGPUErrors(mcblasSetMathMode(handle, MCBLAS_TENSOR_OP_MATH));

  /* Time the "optimized" implementation */
  mcEvent_t start, stop;
  // Allocate MACA events that we'll use for timing
  checkGPUErrors(mcEventCreate(&start));
  checkGPUErrors(mcEventCreate(&stop));

  // printf( "create Handle\n");
    /* Allocate space for the matrices */
    /* Note: I create an extra column in A to make sure that
       prefetching beyond the matrix does not cause a segfault */
    // const size_t mem_size_A = m * (k + 1) * sizeof(float);
    const size_t mem_size_A = m * k * sizeof(float);
    const size_t mem_size_B = k * n * sizeof(float);
    const size_t mem_size_C = m * n * sizeof(float);
    a = (float *)malloc(mem_size_A);
    b = (float *)malloc(mem_size_B);
    c = (float *)malloc(mem_size_C);
    cold = (float *)malloc(mem_size_C);
    cref = (float *)malloc(mem_size_C);

    /* Generate random matrices A, B, Cold */
    random_matrix(m, k, a);
    random_matrix(k, n, b);
    memset(cold, 0, mem_size_C);
    memset(cref, 0, mem_size_C);

    /* Init device matrix*/
    float *d_A, *d_B, *d_C, *d_Cref;
    checkGPUErrors(mcMalloc((void **)&d_A, mem_size_A));
    checkGPUErrors(mcMalloc((void **)&d_B, mem_size_B));
    checkGPUErrors(mcMemcpy(d_A, a, mem_size_A, mcMemcpyHostToDevice));
    checkGPUErrors(mcMemcpy(d_B, b, mem_size_B, mcMemcpyHostToDevice));
    checkGPUErrors(mcMalloc((void **)&d_C, mem_size_C));
    checkGPUErrors(mcMalloc((void **)&d_Cref, mem_size_C));

    /* Run the reference implementation so the answers can be compared */
    if (!trace) {
    // only run the reference implementation when not in trace mode
      SGEMM_GPU_REF(handle, m, n, k, d_A, d_B, d_Cref);
      checkGPUErrors(mcMemcpy(cref, d_Cref, mem_size_C, mcMemcpyDeviceToHost));
      checkGPUErrors(mcDeviceSynchronize());

      // Warp up 
      SGEMM_GPU(handle, m, n, k, d_A, d_B, d_C);
      checkGPUErrors(mcDeviceSynchronize());
      checkGPULastErrors();
      // copy result from device to host
      checkGPUErrors(mcMemcpy(cold, d_C, mem_size_C, mcMemcpyDeviceToHost));

      diff = compare_matrices(m, n, cold, cref);
      if (diff > 0.5f || diff < -0.5f) {
        printf("diff too big !\n");
        exit(-1);
      }
    }
    // Record the start event
    checkGPUErrors(mcEventRecord(start, NULL));

    for (int rep = 0; rep < num_rep; rep++) {
      /* Time your implementation */
      SGEMM_GPU(handle, m, n, k, d_A, d_B, d_C);
    }
    // Record the stop event
    checkGPUErrors(mcEventRecord(stop, NULL));
    // Wait for the stop event to complete
    checkGPUErrors(mcEventSynchronize(stop));
    float msecTotal = 0.0f;
    checkGPUErrors(mcEventElapsedTime(&msecTotal, start, stop));

    // Compute and print the performance
    float msecPerMatrixMul = msecTotal / num_rep;
    double flopsPerMatrixMul = 2.0 * m * k * n;
    double gflops =
        (flopsPerMatrixMul * 1.0e-9f) / (msecPerMatrixMul / 1000.0f);

    printf("%.2f,%le\n", gflops, diff);
    free(a);
    free(b);
    free(c);
    free(cold);
    free(cref);

    checkGPUErrors(mcFree(d_A));
    checkGPUErrors(mcFree(d_B));
    checkGPUErrors(mcFree(d_C));
    checkGPUErrors(mcFree(d_Cref));

  // Destroy the handle
  checkGPUErrors(mcblasDestroy(handle));
}

int main(int argc, char **argv) {
  // parse CLI
  int opt;
  int m, n, k; // problem size
  bool sweep = false, trace = false;
  // get and print gpu info
  mcDeviceProp_t deviceProp;
  int devID = 0;
  checkGPUErrors(mcSetDevice(devID));
  auto error = mcGetDeviceProperties(&deviceProp, devID);
  if (error != mcSuccess) {
    printf("mcGetDeviceProperties returned error code %d, line(%d)\n", error,
           __LINE__);
    exit(EXIT_FAILURE);
  }
  
  while ((opt = getopt(argc, argv, "ht")) != -1) {
    switch (opt) {
        case 'h':
            print_help(argv[0]);
            return 0;
        case 't':
            trace = true;
            break;
        default:
            print_help(argv[0]);
            return 1;
    }
  }
  // Check for remaining arguments for M, N, K
  if (optind < argc) {
      int remaining_args = argc - optind;
      if (remaining_args == 3) {
          m = std::atoi(argv[optind]);
          n = std::atoi(argv[optind + 1]);
          k = std::atoi(argv[optind + 2]);
          callKernel(m, n, k, trace);
          printf("GPU Device %d: \"%s\" with compute capability %d.%d\n", devID,
              deviceProp.name, deviceProp.major, deviceProp.minor);
      } else {
          print_help(argv[0]);
          return 1;
      }
  } else {
      // sweep = true;
      printf("MatrixSize,GFLOPS,Diff\n");
      for (int p = kStart; p <= kStop; p += kStep) {
        printf("%d,", p);
        callKernel(p, p, p, trace);
      }
  }

  return 0;
}
